package main

import (
	"fmt"
	"math/rand"
	"sync"
	"time"
)

type SlidingWindowCounter struct {
	mutex *sync.RWMutex // 读写锁

	buckets    map[int64]int64 // 桶，key 为时间戳（秒），val 为对应时间点的计数数量
	bucketSize int64           // 每个桶的大小限制
	windowSize int64           // 窗口大小，单位：秒
}

func NewSlidingWindowCounter(windowSize, bucketSize int64) *SlidingWindowCounter {
	return &SlidingWindowCounter{
		mutex: &sync.RWMutex{},

		buckets:    make(map[int64]int64),
		bucketSize: bucketSize,
		windowSize: windowSize,
	}
}

// removeOldBuckets 移除过期的桶。
func (s *SlidingWindowCounter) removeOldBuckets() {
	// 当前时间减去桶长度
	now := time.Now().Unix() - s.windowSize
	for timestamp := range s.buckets {
		// 移除保留的桶中过期部分
		if timestamp <= now {
			delete(s.buckets, timestamp)
		}
	}
}

// Increment 在最新的桶中增加计数。
func (s *SlidingWindowCounter) Increment(i int64) {
	if i == 0 {
		return
	}

	s.mutex.Lock()
	defer s.mutex.Unlock()

	// 拿到当前时间戳
	now := time.Now().Unix()
	// 如果增加计数后当前桶容量会溢出，则更新当前桶计数至最大容量（即舍弃溢出的计数）
	// 否则更新桶计数
	if s.buckets[now]+i > s.bucketSize {
		s.buckets[now] = s.bucketSize
	} else {
		s.buckets[now] += i
	}
	// 清除过期的桶
	s.removeOldBuckets()
}

// Sum 统计当前窗口内的计数数量。
func (s SlidingWindowCounter) Sum(now time.Time) int64 {
	s.mutex.RLock()
	defer s.mutex.RUnlock()

	sum := int64(0)
	for timestamp, value := range s.buckets {
		// 统计合法时间内的桶的计数数量
		if timestamp >= now.Unix()-s.windowSize {
			sum += value
		}
	}

	return sum
}

func main() {
	s := NewSlidingWindowCounter(10, 10)
	go func() {
		for range time.NewTicker(time.Second).C {
			s.Increment(rand.Int63n(10))
			fmt.Println(s.buckets, s.Sum(time.Now()))
		}
	}()

	go func() {
		for range time.NewTicker(time.Second).C {
			s.Increment(rand.Int63n(10))
			fmt.Println(s.buckets, s.Sum(time.Now()))
		}
	}()

	select {}
}
